import numpy as np
import cv2
import datetime
import base64
from core import db, logger
from modules.image import ImageModel

class CollectionService():                 
    def cmpHash(self, hash1, hash2):
        n = 0
        if len(hash1) != len(hash2):
            return -1
        for i in range(len(hash1)):
            if hash1[i] != hash2[i]:
                n = n + 1
        print(f"cmpHash: {hash1} {hash2} {n}")
        return n
    
    
    def hash(self, img):
        hash_str = ''
        img = cv2.resize(img, (9, 8))
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        for i in range(8):
            for j in range(8):
                if gray[i, j] > gray[i, j+1]:
                    hash_str = hash_str+'1'
                else:
                    hash_str = hash_str+'0'
        return hash_str

    def is_duplicate(self, hash_val: str) -> bool:
        for h in self.seen_frames:
            if self.cmpHash(h, hash_val) < 15:
                return True
        return False

    def add_frame(self, output_dir: str, ds_id: int, canmera_id: int, frame: np.ndarray):
        hash_val = self.hash(frame)
        if hash_val is not None and not self.is_duplicate(hash_val):
            self.seen_frames.add(hash_val)  # 添加到已保存集合中
            file_name = datetime.datetime.now().strftime('%Y-%m-%d-%H%M%S') + '.jpg'
            # 确保输出目录以斜杠结尾
            if not output_dir.endswith('/'):
                output_dir += '/'
            # 保存图片文件
            cv2.imwrite(output_dir + file_name, frame, [cv2.IMWRITE_JPEG_QUALITY, 90])

            img = ImageModel()
            img.name = file_name
            img.ds_id = ds_id
            img.camera_id = canmera_id
            img.hash = hash_val
            img.type = 'train'
            img.status = '0'
            db.session.add(img)
            db.session.commit()
            logger.info(f"Frame with hash {hash_val} add to {ds_id} : {canmera_id}")
            return True
        else:
            logger.info(f"Frame with hash {hash_val} already exists.")
            return False

    def load_frame(self, ds_id: int, canmera_id: int):
        self.seen_frames = set()
        frames = ImageModel.query.filter_by(ds_id=ds_id, camera_id=canmera_id).all()
        for frame in frames:
            if frame.hash:
                self.seen_frames.add(frame.hash)
